from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import ASGIApp
from starlette import status
import uvicorn
from config import settings
from config.handler import start_app_handler, stop_app_handler
from api.router import api_router
from api.storage import router as ossRouter
from api.core import router as coreRouter
import argparse

# 限制上传最大数量
# @see: https://github.com/fastapi/fastapi/discussions/8167
class LimitUploadSize(BaseHTTPMiddleware):
    def __init__(self, app: ASGIApp, max_upload_size: int) -> None:
        super().__init__(app)
        self.max_upload_size = max_upload_size

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        if request.method == 'POST':
            if 'content-length' not in request.headers:
                return Response(status_code=status.HTTP_411_LENGTH_REQUIRED)
            content_length = int(request.headers['content-length'])
            if content_length > self.max_upload_size:
                return Response(status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE)
        return await call_next(request)
    
def get_app(args:argparse.Namespace) -> FastAPI:
    fast_app = FastAPI(title=settings.PROJECT_NAME, debug=settings.IS_DEBUG)
    fast_app.include_router(api_router, prefix="/api")
    fast_app.include_router(ossRouter, prefix="/oss")
    fast_app.include_router(coreRouter, prefix="/auth")
    fast_app.add_event_handler("startup", start_app_handler(fast_app, args))
    fast_app.add_event_handler("shutdown", stop_app_handler(fast_app))
    origins = settings.CORS_ORIGINS.split(',')
    fast_app.add_middleware(
        CORSMiddleware,
        allow_origins=origins,
        allow_credentials=True,
        allow_methods=["DELETE", "GET", "POST", "PUT"],
        allow_headers=["*"],
    )
    fast_app.add_middleware(LimitUploadSize, max_upload_size=settings.MAX_UPLOAD_SIZE)
    return fast_app

def load_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--nlp", default="nllb", help="翻译模型")
    parser.add_argument("--tts", default="mms", help="文本生成语音模型",)
    parser.add_argument("--asr", default="faster-whisper", help="语音识别模型", )
    return parser.parse_args()


if __name__ == "__main__":
    # setup_log()
    args = load_args()
    app = get_app(args)
    uvicorn.run(app, host=settings.HOST, port=settings.PORT, log_config='./config/logging.yaml') # './config/logging.yaml'
